package com.reebake.ideal.servlet.core;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;

import java.util.*;

/**
 * 请求包装，控制header
 */
public class HttpHeaderRequestWrapper extends HttpServletRequestWrapper {
    private final Map<String, String> customHeaders = new HashMap<>();

    public HttpHeaderRequestWrapper(HttpServletRequest request, String headerName, String headerValue) {
        super(request);
        this.customHeaders.put(headerName, headerValue);
    }

    public HttpHeaderRequestWrapper(HttpServletRequest request, Map<String, String> headers) {
        super(request);
        customHeaders.putAll(headers);
    }

    @Override
    public String getHeader(String name) {
        if(!customHeaders.containsKey(name)) {
            return super.getHeader(name);
        }
        return customHeaders.get(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        Set<String> headerNames = new HashSet<>(customHeaders.keySet());
        headerNames.addAll(Collections.list(super.getHeaderNames()));
        return Collections.enumeration(headerNames);
    }
}